// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <map>
#include <unordered_map>
#include <kfor/sorted_array.h>

namespace kfor {

    typedef uint32_t last_id_t;

    /*
        Compressed chain of blocks that store the document IDs and offsets of a given token.
        Offsets of singular and multi-valued fields are encoded differently.
    */
    class IDList {
    public:

        // A block stores a sorted list of Document IDs compactly
        struct Block {
            SortedArray ids;

            // link to next block
            Block *next = nullptr;

            bool contains(uint32_t id);

            uint32_t upsert(uint32_t id);

            uint32_t erase(uint32_t id);

            uint32_t size() {
                return ids.getLength();
            }
        };

        class Iterator {
        private:
            Block *curr_block;
            int64_t curr_index;

            Block *end_block;
            std::map<last_id_t, Block *> *id_block_map;

            bool reverse;

        public:
            // uncompressed data structure for performance
            uint32_t *ids = nullptr;

            explicit Iterator(Block *start, Block *end, std::map<last_id_t, Block *> *id_block_map,
                                bool reverse);

            Iterator(Iterator &&rhs) noexcept;

            ~Iterator();

            Iterator &operator=(Iterator &&obj) noexcept;

            [[nodiscard]] bool valid() const;

            void next();

            void previous();

            [[nodiscard]] uint32_t last_block_id() const;

            void skip_n(uint32_t n);

            void skip_to(uint32_t id);

            void reset_cache();

            [[nodiscard]] uint32_t id() const;

            [[nodiscard]] inline uint32_t index() const;

            [[nodiscard]] inline Block *block() const;
        };

        struct ResultIterState {
            const uint32_t *excluded_result_ids = nullptr;
            const size_t excluded_result_ids_size = 0;

            const uint32_t *filter_ids = nullptr;
            const size_t filter_ids_length = 0;

            size_t excluded_result_ids_index = 0;
            size_t filter_ids_index = 0;
            size_t index = 0;

            ResultIterState() = default;

            ResultIterState(uint32_t *excluded_result_ids, size_t excluded_result_ids_size,
                                const uint32_t *filter_ids, const size_t filter_ids_length) : excluded_result_ids(
                    excluded_result_ids),
                                                                                              excluded_result_ids_size(
                                                                                                      excluded_result_ids_size),
                                                                                              filter_ids(filter_ids),
                                                                                              filter_ids_length(
                                                                                                      filter_ids_length) {}
        };

    private:

        // maximum number of IDs (and associated offsets) to store in each block before another block is created
        const uint16_t BLOCK_MAX_ELEMENTS;
        uint32_t ids_length = 0;

        Block root_block;

        // keeps track of the *last* ID in each block and is used for partial random access
        // e.g. 0..[9], 10..[19], 20..[29]
        // MUST be ordered
        std::map<last_id_t, Block *> id_block_map;

        static bool at_end(const std::vector<IDList::Iterator> &its);

        static bool at_end2(const std::vector<IDList::Iterator> &its);

        static bool equals(std::vector<IDList::Iterator> &its);

        static bool equals2(std::vector<IDList::Iterator> &its);

        static void advance_all(std::vector<IDList::Iterator> &its);

        static void advance_all2(std::vector<IDList::Iterator> &its);

        static void advance_non_largest(std::vector<IDList::Iterator> &its);

        static void advance_non_largest2(std::vector<IDList::Iterator> &its);

        static uint32_t advance_smallest(std::vector<IDList::Iterator> &its);

        static uint32_t advance_smallest2(std::vector<IDList::Iterator> &its);

    public:

        explicit IDList(uint16_t max_block_elements);

        ~IDList();

        static void split_block(Block *src_block, Block *dst_block);

        static void merge_adjacent_blocks(Block *block1, Block *block2, size_t num_block2_ids_to_move);

        void upsert(uint32_t id);

        void erase(uint32_t id);

        Block *get_root();

        size_t num_blocks() const;

        size_t num_ids() const;

        uint32_t first_id();

        uint32_t last_id();

        Block *block_of(uint32_t id);

        bool contains(uint32_t id);

        bool contains_atleast_one(const uint32_t *target_ids, size_t target_ids_size);

        Iterator new_iterator(Block *start_block = nullptr, Block *end_block = nullptr);

        Iterator new_rev_iterator();

        static void merge(const std::vector<IDList *> &id_lists, std::vector<uint32_t> &result_ids);

        static void intersect(const std::vector<IDList *> &id_lists, std::vector<uint32_t> &result_ids);

        static bool take_id(ResultIterState &istate, uint32_t id);

        template<class T>
        static bool block_intersect(
                std::vector<IDList::Iterator> &its,
                ResultIterState &istate,
                T func
        );

        uint32_t *uncompress();

        void uncompress(std::vector<uint32_t> &data);

        size_t intersect_count(const uint32_t *res_ids, size_t res_ids_len,
                               bool estimate_facets, size_t facet_sample_interval);
    };

    template<class T>
    bool IDList::block_intersect(std::vector<IDList::Iterator> &its, ResultIterState &istate, T func) {

        switch (its.size()) {
            case 0:
                break;
            case 1:
                while (its[0].valid()) {
                    if (IDList::take_id(istate, its[0].id())) {
                        func(its[0].id(), its, istate.index);
                    }

                    its[0].next();
                }
                break;
            case 2:
                while (!at_end2(its)) {
                    if (equals2(its)) {
                        if (IDList::take_id(istate, its[0].id())) {
                            func(its[0].id(), its, istate.index);
                        }

                        advance_all2(its);
                    } else {
                        advance_non_largest2(its);
                    }
                }
                break;
            default:
                while (!at_end(its)) {
                    if (equals(its)) {
                        //LOG(INFO) << its[0].id();
                        if (IDList::take_id(istate, its[0].id())) {
                            func(its[0].id(), its, istate.index);
                        }

                        advance_all(its);
                    } else {
                        advance_non_largest(its);
                    }
                }
        }

        return false;
    }
}  // namespace kfor
